import Stack from "@/sources/Stack/useSl";
import { Tree } from "../base";
import { RBNode, Color } from "./types";
import { black, color, colorOf, isBlack, isRed, red } from "./util";

/**
 * 二叉搜索树
 */
export default class RBT<T = number> extends Tree<T> {
    private sizes = 0;
    private root: RBNode<T> | null
    // 默认的比较器
    private comparator = (a: T, b: T) => ((a as any) - (b as any))
    constructor(comparator?: (a: T, b: T) => number) {
        super();
        this.root = null
        comparator && (this.comparator = comparator);
    }

    /**
     * 获取根节点
     */
    getRoot() {
        return this.root;
    }

    /**
     * 节点的数量
     */
    size() {
        return this.sizes;
    }

    /**
     * 是否为空
     */
    isEmpty() {
        return this.sizes === 0
    }

    /**
     * 清空二叉树
     */
    clear() {
        // 清空根节点和sizes置为空即可
        this.sizes = 0;
        this.root = null
    }

    /**
     * 添加节点
     * @param element 
     */
    add(element: T) {
        // 首次添加
        if (this.root === null) {
            this.root = new RBNode(element, null, null, null);
            // 根节点，染成黑色
            this.root.color = Color.BLACK
            this.sizes++;
            return;
        }

        // 找到要插入的节点
        let root = this.root;
        let prevNode: RBNode<T> | null,
            isRight = false;
        while (root) {
            // 保存这次的节点
            prevNode = root;
            if (this.comparator(element, root.value!) > 0) {
                // 往右边找
                root = root.right!;
                isRight = true
            } else if (this.comparator(element, root.value!) < 0) {
                // 往左边找
                root = root.left!
                isRight = false;
            } else {
                // 相等，直接更新节点的值
                root.value = element
                // 直接返回
                return;
            }
        }

        // 到这prevNode就是要找的插入的节点

        const newNode = new RBNode(element, prevNode!, null, null)

        // 判断插入左还是右
        if (isRight) {
            // 插入 prevNode的右边
            prevNode!.right = newNode
        } else {
            // 插入prevNode的左边
            prevNode!.left = newNode
        }
        // debugger
        this.afterAdd(newNode)


        this.sizes++;
    }

    /**
     * 添加节点之后修复红黑树，默认插入的节点都是红色的
     * @param node 
     */
    afterAdd(node: RBNode<T>) {

        // 父节点为空，染成黑色（递归到根节点了）
        if (node.parent === null) return black(node)

        // 父节点是黑色的，不用处理
        if (isBlack(node.parent)) return;

        // 父节点是红色的 
        const parent = node.parent,
            grand = node.parent.parent!;
        // 所有的条件都要将grand变为红色的   
        red(grand)
        // 叔叔节点不是红色的
        if (isBlack(node.uncle()!)) {
            // 3. 根据父节点在爷爷节点的左边还是右边， 对爷爷和父节点节点进行旋转
            if (parent.isRight()) { // R
                if (node.isRight()) { // RR
                    // 父节点染成黑色的
                    black(parent)

                    // 左旋转
                    this.rotateLeft(grand)
                } else { // RL
                    // 先进行右旋，再左旋
                    // debugger
                    // 将当前节点染成黑色
                    black(node)
                    // 右旋
                    this.rotateRight(parent)
                    // 左旋
                    this.rotateLeft(grand!)
                }
            } else { // L
                if (node.isLeft()) {  // LL
                    // 父节点染成黑色的
                    black(parent)

                    // LL 右旋转
                    this.rotateRight(grand)
                } else { // LR
                    // 先左旋，再右旋

                    // 将当前节点染成黑色
                    black(node)
                    // 左旋
                    this.rotateLeft(parent)
                    // 右旋, 此时node已经是之前node.parent了
                    this.rotateRight(grand)
                }
            }
        } else { // 上溢
            // 所有的LL, RR, LR, RL上溢结果都一样的
            // 叔叔节点是红色的
            // 将父节点和叔叔节点染成黑色
            black(parent)
            black(node.uncle()!)

            // 将爷爷节点当作新添加之后的处理
            this.afterAdd(grand)
        }
    }

    /**
     * 处理旋转之后父节点的指向
     * @param grand 
     * @param par 
     */
    private afterRotate(grand: RBNode<T>, par: RBNode<T>) {
        // grand原来是左孩子还是右孩子， 先保存起来，后面判断会失误(grand的父亲已经改了)
        const isLeftChild = grand.isLeft();
        // 更新parent
        par.parent = grand.parent
        grand.parent = par;

        // 更改parent的指向
        if (par.parent === null) {
            this.root = par
        } else {
            // 原来grand的父元素不为空
            if (isLeftChild) {
                // 原来是左孩子
                par.parent.left = par
            } else {
                // 原来是右孩子，更改右孩子
                par.parent.right = par
            }
        }
    }

    /**
     * 右旋
     * @param grand 
     * @param par 
     */
    rotateRight(grand: RBNode<T>) {
        let par = grand.left!;

        grand.left = par.right;
        par.right = grand;

        // 更新par.right的parent, 现在已经变成grand.left
        if (grand.left !== null) {
            grand.left.parent = grand
        }
        // 处理旋转之后
        this.afterRotate(grand, par)
    }

    /**
     * 左旋
     * @param grand 
     * @param par 
     */
    rotateLeft(grand: RBNode<T>) {
        const par = grand.right!;

        grand.right = par.left;
        par.left = grand;

        // 更新par.left的parent
        if (grand.right !== null) {
            grand.right.parent = grand;
        }

        // 处理旋转之后
        this.afterRotate(grand, par)
    }



    /**
     * 获取element
     * @param element 
     */
    get(element: T) {
        let root = this.root!;
        while (root !== null) {
            if (this.comparator(element, root.value!) > 0) {
                root = root.right!;
            } else if (this.comparator(element, root.value!) < 0) {
                root = root.left!;
            } else {
                // 找到了
                return root
            }
        }
        // 循环下来都没有找到
        return null
    }

    /**
     * 删除节点
     * @param element 
     */
    remove(element: T) {
        let node = this.get(element);
        // 要删除的节点不存在
        if (node === null) return null

        // 删除度为2的节点, 找到其前驱或者后继节点， 这里使用前驱节点
        if (node.hasTwoChild()) {
            // 1. 获取node节点的前驱节点
            const accesssor = this.accessor(node);

            // 2. 使用前驱节点替换要删除的该节点
            node.value = accesssor!.value

            // 3. 删除该前驱节点, 跟删除度为1和度为2的节点一样
            node = accesssor
        }

        // 判断该节点是父元素的左孩子还是右孩子
        const isLeft = node!.isLeft() ? true : false
        // 删除度为1的节点
        if (node?.hasOneChild()) {
            // 节点的孩子是左还是右
            const hasLeftChild = node.hasLeftChild();
            // 找到替换的节点
            const replacement = (hasLeftChild ? node.left : node.right)!;

            // 该节点是跟节点
            if (node.parent == null) {
                this.root = replacement
            } else {
                if (isLeft) {
                    node.parent!.left = replacement
                } else {
                    node.parent!.right = replacement
                }
            }
            
            // 更新父节点
            replacement.parent = node.parent
            this.afterRemove(node, replacement)
        } else if (node!.parent === null) {
            // 删除的是根节点
            this.root = null
        }
        else {
            // 度为0的节点
            if (isLeft) {
                // 是父节点的左孩子
                node!.parent!.left = null
            } else {
                // 是父节点的右孩子
                node!.parent!.right = null
            }
            this.afterRemove(node!, null)
        }
        // 长度减一
        this.sizes--;

    }

    /**
     * 红黑树中删除的都是叶子节点, 删除一度或者二度的节点,最终都是找到替换的节点将其删除了
     * 删除节点之后调节平衡
     */
    afterRemove(node: RBNode<T>, replacement: RBNode<T> | null) {
        // debugger
        // 1. 删除的节点是红色的，不用处理
        if (isRed(node)) return;

        // 2. 替换节点是红色的，将替代的节点染成黑色的
        if (isRed(replacement)) {
            black(replacement)
            return;
        }

        // 根节点
        if (node.parent === null) return;

        // 删除的节点必定是叶子节点

        const parent = node.parent!;
        // 判断删除的节点在父元素的左边还是右边
        // const isRemoveLeft = parent.left ? false : true

        // 如果左子节点不为空，就看它是不是父亲的左子节点
        // node.isLeft() 兼容递归afterRemove的情况
        const isRemoveLeft = parent.left === null || node.isLeft()

        // 获取兄弟节点 (如果左边被删除了，那么兄弟节点就是右边，反之)
        let sibling = (isRemoveLeft ? parent.right : parent.left)!;
        if (isRemoveLeft) {
            // 删除的节点在左边 (跟右边是对称的)， 兄弟节点在右边

            // 兄弟节点是红色 (说明在b树中，兄弟节点在上方，只能借它的子元素，让其子元素变成删除节点的兄弟节点，变成了处理兄弟节点是黑色的情况，后面统一处理)
            if (isRed(sibling)) {
                // 1.兄弟节点染成黑色，父节点染成红色
                black(sibling);
                red(parent);

                // 2. parent左旋转
                this.rotateLeft(parent);

                //! 旋转之后记得更改sibling 
                sibling = parent.right!;
            }


            // 兄弟节点是黑色的情况
            // 1. 兄弟节点至少有一个红色子节点 (跟兄弟节点借)
            // 2. 兄弟节点没有红色子节点 (父元素下来合并)

            if (isBlack(sibling!.right) && isBlack(sibling!.left)) {
                // 兄弟没有红色节点，父节点下来合并

                const isParentBlack = isBlack(parent)
                // 1. 将父节点染成黑色，兄弟节点染成红色
                black(parent);
                red(sibling)

                // 2. 原来父节点是黑色的，下来后也可能倒置下溢，将父节点当作被删除的元素处理
                if (isParentBlack) {
                    this.afterRemove(parent, null)
                }
            } else {
                // 至少有一个红色子节点，借兄弟的红色节点

                // 分为 RR, RL,

                // 先处理RL的情况先右旋，后面统一左左旋

                if (isBlack(sibling.right)) {
                    // 右边为黑色的，左边是红色的，说明是RL的情况

                    // 1. sibling右旋转
                    this.rotateRight(sibling!)

                    // 2. 旋转之后sibling就变为 原来的sibling的右孩子
                    sibling = parent.right!
                }

                // sibling继承父节点的颜色
                // 染成父节点的颜色
                color(sibling, colorOf(parent))

                // 旋转之后两边变为黑色
                black(parent);
                black(sibling.right)

                // 对parent左旋转
                this.rotateLeft(parent);
            }
        } else {
            // 删除的节点在右边

            // 兄弟节点是红色 (说明在b树中，兄弟节点在上方，只能借它的子元素，让其子元素变成删除节点的兄弟节点，变成了处理兄弟节点是黑色的情况，后面统一处理)
            if (isRed(sibling)) {
                // 1.兄弟节点染成黑色，父节点染成红色
                black(sibling);
                red(parent);

                // 2. parent右旋转
                this.rotateRight(parent);

                //! 旋转之后记得更改sibling
                sibling = parent.left!
            }


            // 兄弟节点是黑色的情况
            // 1. 兄弟节点至少有一个红色子节点 (跟兄弟节点借)
            // 2. 兄弟节点没有红色子节点 (父元素下来合并)

            if (isBlack(sibling!.right) && isBlack(sibling!.left)) {
                // 兄弟没有红色节点

                const isParentBlack = isBlack(parent)
                // 1. 将父节点染成黑色，兄弟节点染成红色
                black(parent);
                red(sibling)

                // 2. 原来父节点是黑色的，下来后也可能倒置下溢，将父节点当作被删除的元素处理
                if (isParentBlack) {
                    this.afterRemove(parent, null)
                }

            } else {
                // 至少有一个红色子节点

                // 分为 LL, LR,

                // 先处理LR的情况先左旋，后面统一左右旋
                if (isBlack(sibling.left)) {
                    // 左孩子为黑色的，右孩子是红，说明是LR的情况

                    // 1. sibling左旋转
                    this.rotateLeft(sibling!)

                    // 2. 旋转之后sibling就变为 原来的sibling的左孩子
                    sibling = parent.left!
                }

                // sibling继承父节点的颜色
                // 染成父节点的颜色
                color(sibling, colorOf(parent))

                // 旋转之后两边变为黑色
                black(parent);
                black(sibling.left)

                // 对parent右旋转
                this.rotateRight(parent);

            }
        }
    }



    /**
     * 后继节点
     * 从根节点的右节点开始，不断找该右节点的左节点
     */
    successor(root: RBNode<T>) {
        // 树是空的
        if (this.isEmpty()) return null;
        let node = root?.right;
        // 没有右节点
        if (node === null) return null;

        while (node!.left !== null) {
            node = node?.left;
        }
        return node;
    }

    /**
     * 前继节点
     * 从根节点的左节点开始，不断找该左节点的右节点
     */
    accessor(root: RBNode<T>) {
        if (this.isEmpty()) return null;
        let node = root?.left;
        // 没有左节点
        if (node === null) return null;

        while (node!.right !== null) {
            node = node?.right;
        }
        return node;
    }




    /**
     * 是否包含元素
     * @param element 
     */
    contains(element: T) {
        if (this.isEmpty()) return false
        return this.get(element) !== null;
    }

    /**
     * 批量添加
     * @param eles 
     */
    batchAdd(eles: Array<T>) {
        eles.forEach(ele => {
            // if (ele as any === 13) debugger
            this.add(ele)
            // print(this)
        })
    }

    /**
     * 批量删除
     * @param elem 
     */
    batchRemove(eles: Array<T>) {
        eles.forEach(ele => {
            this.remove(ele)
            // print(this)
        })
    }


    /**
      * 中序遍历(非递归)
      * @param node 
      */
    inOrder(iterator: (ele: T) => void, node = this.root) {
        if (node === null) return;
        // 创建一个栈
        const stack = new Stack<RBNode<T>>()

        // 1. 开始时node为根节点，进入循环
        // 2. 最左节点为空，栈中还有值
        // 3. 栈为空了，开始遍历右边
        while (node === this.root || !stack.isEmpty() || node !== null) {
            // 从node开始一直遍历到最左边
            while (node !== null) {
                stack.push(node);
                node = node.left;
            }
            // 最左边为空， 开始取出最左边的节点开始访问
            node = stack.pop()!;
            // 访问节点
            iterator(node.value!)
            // 访问节点的右节点
            node = node.right;
        }
    }
} 